import mysql.connector
import hashlib
import os
import sys
from mysql.connector import errorcode

# 数据库配置
DB_CONFIG = {
    'user': 'root',  # MySQL用户名
    'password': '123456',  # MySQL密码
    'host': 'localhost',  # MySQL服务器地址
    'database': '3ftp_db'  # 数据库名称
}


def create_database(cursor):
    """创建数据库"""
    try:
        cursor.execute(
            "CREATE DATABASE IF NOT EXISTS {} DEFAULT CHARACTER SET 'utf8mb4'".format(DB_CONFIG['database']))
        print(f"数据库 {DB_CONFIG['database']} 创建成功或已存在")
    except mysql.connector.Error as err:
        print(f"创建数据库失败: {err}")
        sys.exit(1)


def create_tables(cnx):
    """创建数据表"""
    cursor = cnx.cursor()

    # 创建用户表 - 添加 public_key 字段
    users_table = (
        "CREATE TABLE IF NOT EXISTS `users` ("
        "  `id` INT AUTO_INCREMENT PRIMARY KEY,"
        "  `username` VARCHAR(50) UNIQUE NOT NULL,"
        "  `password_hash` VARCHAR(64) NOT NULL,"  # SHA256结果长度为64
        "  `salt` VARCHAR(32) NOT NULL,"  # 盐值，16字节的十六进制表示（32字符）
        "  `home_dir` VARCHAR(255) NOT NULL,"
        "  `public_key` TEXT"  # 添加公钥字段，初始为NULL，用户首次登录时更新
        ") ENGINE=InnoDB"
    )

    try:
        print("创建用户表...")
        cursor.execute(users_table)
    except mysql.connector.Error as err:
        if err.errno == errorcode.ER_TABLE_EXISTS_ERROR:
            print("用户表已存在")
        else:
            print(f"创建用户表失败: {err}")
            sys.exit(1)

    # 添加默认用户
    default_users = [
        ('zhangsan', 'zhangsan1234', '/home/ftp/zhangsan'),
        ('lisi', 'lisi567', '/home/ftp/lisi'),
        ('wangwu', 'wangwu123', '/home/ftp/wangwu'),
        ('liuyun', 'liuyun456', '/home/ftp/liuyun'),
        ('liuchang', 'liuchang789', '/home/ftp/liuchang'),
        ('wanghong', 'wanghong001', '/home/ftp/wanghong'),
        ('zhaoliu', 'zhaoliu2023', '/home/ftp/zhaoliu'),
        ('chenxi', 'chenxi888', '/home/ftp/chenxi'),
        ('sunming', 'sunming999', '/home/ftp/sunming'),
        ('zhoujie', 'zhoujie666', '/home/ftp/zhoujie')
    ]

    for username, password, home_dir in default_users:
        # 检查用户是否已存在
        cursor.execute("SELECT 1 FROM users WHERE username = %s", (username,))
        if cursor.fetchone():
            continue

        # 生成随机盐
        salt = os.urandom(16).hex()
        # 使用SHA-256哈希密码
        password_hash = hashlib.sha256((password + salt).encode()).hexdigest()
        
        # 创建用户主目录
        if not os.path.exists(home_dir):
            os.makedirs(home_dir, exist_ok=True)

        # 插入用户（不包含公钥，设置为NULL）
        add_user = ("INSERT INTO users "
                    "(username, password_hash, salt, home_dir, public_key) "
                    "VALUES (%s, %s, %s, %s, NULL)")

        try:
            cursor.execute(add_user, (username, password_hash, salt, home_dir))
            print(f"添加用户: {username} (公钥将在首次登录时生成)")
        except mysql.connector.Error as err:
            print(f"添加用户 {username} 失败: {err}")

    cnx.commit()
    cursor.close()


def init_database():
    """初始化数据库"""
    try:
        # 先连接但不指定数据库，用于创建数据库
        cnx = mysql.connector.connect(
            user=DB_CONFIG['user'],
            password=DB_CONFIG['password'],
            host=DB_CONFIG['host']
        )
        cursor = cnx.cursor()

        # 创建数据库
        create_database(cursor)
        cursor.close()
        cnx.close()

        # 重新连接指定数据库
        cnx = mysql.connector.connect(**DB_CONFIG)
        print(f"成功连接到数据库 {DB_CONFIG['database']}")

        # 创建表并添加默认用户
        create_tables(cnx)

        # 验证表创建
        cursor = cnx.cursor()
        cursor.execute("SHOW TABLES")
        tables = cursor.fetchall()

        # 显示用户数据
        cursor.execute("SELECT username, home_dir FROM users")
        users = cursor.fetchall()

        cursor.close()
        cnx.close()

        print("\n数据库初始化完成!")
        print("用户将在首次登录时自动生成RSA密钥对")

    except mysql.connector.Error as err:
        print(f"数据库错误: {err}")
        sys.exit(1)
    except Exception as e:
        print(f"初始化错误: {e}")
        sys.exit(1)


if __name__ == '__main__':
    # 修改为您的MySQL root密码
    DB_CONFIG['password'] = input("请输入MySQL root密码: ")
    init_database()